import sys
sys.path.append("/Users/ender/Code_repo/rosalind-brush-questions")
from utils import read_fasta


class NucleotideBase:
    def __init__(self, nb: str) -> None:
        self.nucleotide_base = nb
    
    def transition(self) -> str:
        if self.nucleotide_base == "A":
            return "G"
        if self.nucleotide_base == "G":
            return "A"
        if self.nucleotide_base == "C":
            return "T"
        if self.nucleotide_base == "T":
            return "C"
    
    def transversion(self) -> str:
        if self.nucleotide_base == "A":
            return ["C", "T"]
        if self.nucleotide_base == "C":
            return ["A", "G"]
        if self.nucleotide_base == "G":
            return ["T", "C"]
        if self.nucleotide_base == "T":
            return ["G", "A"]
        
    def is_transition(self, new_nb: str) -> bool:
        return self.transition() == new_nb
    
    def is_tranversion(self, new_nb: str) -> bool:
        return new_nb in self.transversion()


if __name__ == '__main__':
    # fasta = read_fasta("test.fa")
    # fasta = read_fasta("14_TransitionsAndTransversions/test.fa")
    fasta = read_fasta("./rosalind_tran.txt")
    s1, s2 = [v for v in fasta.values()]
    n_transition = 0
    n_transversion = 0
    for nb1, nb2 in zip(s1, s2):
        nb = NucleotideBase(nb1)
        n_transition += nb.is_transition(nb2)
        n_transversion += nb.is_tranversion(nb2)
    
    print(n_transition / n_transversion)
